import os
import logging
import click
import shutil

from plotting.plot_tools import load_plots
from plotting.byte_types import hexstr_to_bytes
from plotting.check_plots import check_plots

from blspy import G1Element, PrivateKey

logging.basicConfig(level=logging.NOTSET)  # 设置日志级
log = logging.getLogger(__name__)


@click.group()
def cli():
    pass


@click.command("version", short_help="Show chia version")
def version_cmd() -> None:
    print("0.1.0")


@click.command("show", short_help="Show plot info")
@click.option('--file', help='plot file name')
def show_cmd(file: str) -> None:
    show_memo = True
    root_path = os.path.dirname(file)
    match_str = os.path.basename(file)

    _, provers, failed_to_open_filenames, no_key_filenames = load_plots(
        {},
        {},
        [],
        [],
        match_str,
        show_memo,
        root_path,
        open_no_key_filenames=True,
    )

    for plot_path, plot_info in provers.items():
        pr = plot_info.prover
        log.info(f"Testing plot {plot_path} k={pr.get_size()}")
        log.info(f"\tPool public key: {plot_info.pool_public_key}")


@click.command("clear", short_help="Clear invalid plot file")
@click.option('--fpk', help='plotting fpk')
@click.option('--ppk', help='plotting ppk')
@click.option('--dir', help='root dir')
def clear_cmd(fpk: str, ppk: str, dir: str):
    match_str = None
    show_memo = True
    root_path = dir

    fpk_key = G1Element.from_bytes(hexstr_to_bytes(fpk))
    ppk_key = G1Element.from_bytes(hexstr_to_bytes(ppk))

    _, provers, failed_to_open_filenames, no_key_filenames = load_plots(
        {},
        {},
        [fpk_key],
        [ppk_key],
        match_str,
        show_memo,
        root_path,
        open_no_key_filenames=False
    )

    log.info(f"")

    if len(failed_to_open_filenames) > 0 or len(no_key_filenames) > 0:

        for plot_path in failed_to_open_filenames:
            log.info(f"Clear invalid plot {plot_path}")

            if os.path.exists(plot_path):
                os.remove(plot_path)

        for plot_path in no_key_filenames:
            log.info(f"Clear no_key plot {plot_path}")

            if os.path.exists(plot_path):
                os.remove(plot_path)

    else:
        log.info(f"All plot valid")


@click.command("check", short_help="Checks plots")
@click.option('--fpk', help='plotting fpk')
@click.option('--ppk', help='plotting ppk')
@click.option('--dir', help='root dir')
@click.option("-n", "--num", help="Number of plots or challenges", type=int, default=None)
@click.option(
    "-g",
    "--grep_string",
    help="Shows only plots that contain the string in the filename or directory name",
    type=str,
    default=None,
)
@click.option("-l", "--list_duplicates", help="List plots with duplicate IDs", default=False, is_flag=True)
@click.option("--debug-show-memo", help="Shows memo to recreate the same exact plot", default=False, is_flag=True)
@click.option("--challenge-start", help="Begins at a different [start] for -n [challenges]", type=int, default=None)
def check_cmd(
    fpk: str, ppk: str, dir: str, num: int, grep_string: str, list_duplicates: bool, debug_show_memo: bool, challenge_start: int
):
    fpk_key = G1Element.from_bytes(hexstr_to_bytes(fpk))
    ppk_key = G1Element.from_bytes(hexstr_to_bytes(ppk))

    check_plots(fpk_key, ppk_key, dir, num, challenge_start, grep_string, list_duplicates, debug_show_memo)


cli.add_command(version_cmd)
cli.add_command(show_cmd)
cli.add_command(clear_cmd)
cli.add_command(check_cmd)


if __name__ == "__main__":
    cli()